package com.dc.project.util.reader;

import com.dc.project.util.reader.base.BaseReaderUtil;
import com.dc.project.util.reader.base.FileReaderModel;
import com.own.util.ImageCompressorUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;

import javax.imageio.ImageIO;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;

/**
 * PdfReaderUtil
 *
 * @author chenxueli
 * @date 2025-03-08 16:12:00
 */
public class PdfReaderUtil implements BaseReaderUtil {

    /**
     * 逐页处理PDF文档内容
     *
     * @param document PDF文档
     * @return 处理后的文本内容
     * @throws IOException 异常
     */
    private static FileReaderModel processDocumentByPage(PDDocument document) throws IOException {
        var fullResult = new StringBuilder();
        var pageCount = document.getNumberOfPages();
        var pdfRenderer = new PDFRenderer(document);
        var fileList = new ArrayList<File>();
        try {
            // 逐页处理
            for (int i = 0; i < pageCount; i++) {
                var stripper = new PDFTextStripper();
                stripper.setSortByPosition(true);
                // 设置仅提取当前页的文本
                stripper.setStartPage(i + 1);
                stripper.setEndPage(i + 1);
                var pageText = stripper.getText(document).trim();
                // 判断当前页是否为空白或几乎没有文字的页面
                if (StringUtils.isBlank(pageText) || pageText.length() < 20) {
                    // 如果文本内容为空或很少，认为是图片页，使用OCR处理
                    var image = pdfRenderer.renderImageWithDPI(i, 300);
                    var tempFile = File.createTempFile("page-" + i, ".png");
                    ImageIO.write(image, "png", tempFile);
                    // 如果是大于2MB的图片，需要将图片进行压缩
                    var maxSize = 2 * 1024;
                    fileList.add(tempFile.length() > (maxSize * 1024)
                            ? ImageCompressorUtil.resizeImageProportionally(tempFile, 2048, 2048, 1)
                            : tempFile
                    );
                    fullResult.append(FileReaderModel.getPlaceholder(fileList.size() - 1));
                } else {
                    // 如果页面包含足够的文本，使用文本提取
                    fullResult.append(pageText);
                }
                fullResult.append("\n\n");
            }
        } catch (Exception e) {
            return FileReaderModel.EMPTY;
        }
        return FileReaderModel.newInstance(fullResult.toString(), fileList);
    }

    /**
     * 读取文件内容
     *
     * @param file 文件
     * @return 文件内容
     */
    @Override
    public FileReaderModel read(File file) {
        try (var document = PDDocument.load(file)) {
            // 按页处理，检查每页是否包含文本
            return processDocumentByPage(document);
        } catch (IOException e) {
            return FileReaderModel.EMPTY;
        }
    }

}
